import csv
import numpy as np
from scipy.stats import norm
from mayavi import mlab

# 气体参数
Wi = 0.4 * 28 + 0.6 * 41  # eV/电荷
Fano = 0.3  # Fano因子
Dt = 0.008  # 横向扩散系数 (mm)
Dl = 0.008  # 纵向扩散系数 (mm)

def read_deposits_info(csv_path):
    deposits = []
    with open(csv_path, 'r') as f:
        reader = csv.reader(f)
        in_deposits = False
        header = None
        for row in reader:
            row = [cell.strip() for cell in row]
            if not row:
                continue
            if "Deposits Info" in row:
                in_deposits = True
                try:
                    header = next(reader)
                except StopIteration:
                    print("错误：未找到表头")
                    return []
                continue
            if in_deposits and header:
                try:
                    pre_x = float(row[header.index('pre_pos_x')])
                    pre_y = float(row[header.index('pre_pos_y')])
                    pre_z = float(row[header.index('pre_pos_z')])
                    post_x = float(row[header.index('post_pos_x')])
                    post_y = float(row[header.index('post_pos_y')])
                    post_z = float(row[header.index('post_pos_z')])
                    edep_MeV = float(row[header.index('edep')])
                    
                    deposits.append({
                        'pre_pos_x': pre_x,
                        'pre_pos_y': pre_y,
                        'pre_pos_z': pre_z,
                        'post_pos_x': post_x,
                        'post_pos_y': post_y,
                        'post_pos_z': post_z,
                        'edep_MeV': edep_MeV
                    })
                except (IndexError, ValueError) as e:
                    print(f"数据解析错误：{e}")
    return deposits

def process_deposits(deposits):
    charge_positions = []
    for idx, dep in enumerate(deposits):
        mid_x = (dep['pre_pos_x'] + dep['post_pos_x']) / 2
        mid_y = (dep['pre_pos_y'] + dep['post_pos_y']) / 2
        mid_z = (dep['pre_pos_z'] + dep['post_pos_z']) / 2
        
        edep_MeV = dep['edep_MeV']
        if edep_MeV <= 0:
            continue
            
        edep_eV = edep_MeV * 1e6
        m_ion = edep_eV / Wi
        s_ion = np.sqrt(Fano * m_ion)
        
        n_ion_float = norm.rvs(loc=m_ion, scale=s_ion)
        n_ion = max(0, int(n_ion_float))
        
        cov = [[Dt**2, 0, 0], [0, Dt**2, 0], [0, 0, Dl**2]]
        
        for _ in range(n_ion):
            noise = np.random.multivariate_normal([0,0,0], cov)
            charge_positions.append([
                mid_x + noise[0],
                mid_y + noise[1],
                mid_z + noise[2]
            ])
    
    return np.array(charge_positions)

def visualize_3d(charge_positions, voxel_size=0.1):
    if len(charge_positions) == 0:
        print("无有效电荷数据")
        return
    
    # 计算空间边界
    min_xyz = np.min(charge_positions, axis=0)
    max_xyz = np.max(charge_positions, axis=0)
    size = max_xyz - min_xyz
    max_dim = np.max(size) * 1.2  # 扩展20%边界
    center = (min_xyz + max_xyz) / 2
    boundaries = np.array([
        center[0] - max_dim/2,
        center[0] + max_dim/2,
        center[1] - max_dim/2,
        center[1] + max_dim/2,
        center[2] - max_dim/2,
        center[2] + max_dim/2
    ])
    
    # 创建三维网格
    x_edges = np.arange(boundaries[0], boundaries[1], voxel_size)
    y_edges = np.arange(boundaries[2], boundaries[3], voxel_size)
    z_edges = np.arange(boundaries[4], boundaries[5], voxel_size)
    
    # 统计体元电荷数
    hist, edges = np.histogramdd(
        charge_positions,
        bins=(x_edges, y_edges, z_edges)
    )
    
    # 调试输出
    print(f"电荷计数统计：")
    print(f"最小值：{hist.min():.1f}, 最大值：{hist.max():.1f}")
    print(f"平均值：{hist.mean():.1f}, 标准差：{hist.std():.1f}")
    print(f"体元数量：{hist.size}")
    
    # 生成网格坐标
    x, y, z = np.meshgrid(
        (x_edges[:-1] + x_edges[1:])/2,
        (y_edges[:-1] + y_edges[1:])/2,
        (z_edges[:-1] + z_edges[1:])/2,
        indexing='ij'
    )
    
    # 创建可视化窗口
    mlab.figure(size=(1200, 900), bgcolor=(1,1,1))
    
    # 等值面可视化（关键参数调整）
    contours = mlab.contour3d(
        x, y, z, hist,
        contours=20,          # 更多等值面
        opacity=0.8,          # 适当透明度
        colormap='viridis',   # 更直观的颜色映射
        vmax=hist.max(),      # 显式设置最大值
        vmin=hist.min()       # 显式设置最小值
    )
    
    # 添加颜色条和坐标轴
    mlab.colorbar(contours, orientation='vertical', title='Charge Density')
    mlab.axes(
        ranges=boundaries.tolist(),
        xlabel='X (mm)', 
        ylabel='Y (mm)', 
        zlabel='Z (mm)'
    )
    
    # 调整视角
    mlab.view(azimuth=30, elevation=60)
    
    # 手动添加光源（替代mlab.lighting）
    # 创建一个本地光源并设置属性
    light = mlab.points3d(0, 0, 0, scale_factor=0.1, color=(1, 1, 1))
    light.actor.property.lighting = True
    light.actor.property.ambient = 0.5
    light.actor.property.diffuse = 0.5
    light.actor.property.specular = 0.5
    
    # 保存并显示
    mlab.savefig('charge_3d.png', magnification=2)
    mlab.show()

if __name__ == "__main__":
    csv_path = "/opt/star-xp-master/build/source/0-Simulation/e-+,0.01,10/event_0.csv"
    deposits = read_deposits_info(csv_path)
    
    if not deposits:
        print("错误：未读取到任何能量沉积信息")
        exit(1)
    
    charge_positions = process_deposits(deposits)
    
    if len(charge_positions) == 0:
        print("所有能量沉积事件生成0电荷（可能edep过小或单位错误）")
    else:
        print(f"总电荷数量：{len(charge_positions)}")
        visualize_3d(charge_positions, voxel_size=0.1)
